# Author: WY
# Date: 2023/5/17 16:28

import torch
import torch.nn as nn
import torch_geometric as PyG


class GCN(torch.nn.Module):
    def __init__(self, feature_dim, gcn_output_size, temp_size_conv1, temp_size_conv2, grad, node_num, dropout):
        super(GCN, self).__init__()
        # 第一层图卷积
        self.conv1 = PyG.nn.GCNConv(
            in_channels=feature_dim,
            out_channels=temp_size_conv1
        )
        # 第二层图卷积
        self.conv2 = PyG.nn.GCNConv(
            in_channels=temp_size_conv1,
            out_channels=temp_size_conv2
        )
        # Dropout
        self.dropout = nn.Dropout(dropout)
        # 多层感知机
        self.MLP = nn.Sequential(
            nn.Linear(temp_size_conv2, (temp_size_conv2 + gcn_output_size) // grad),
            nn.ReLU(inplace=True),
            nn.Linear((temp_size_conv2 + gcn_output_size) // grad, gcn_output_size)
        )
        # 全连接层(用于结合各站点之间的空间关系)
        self.FC1 = nn.Linear(node_num, 1)
        # ReLU
        self.relu1 = nn.ReLU()

    def forward(self, input, edge_index, edge_weight):
        # 图卷积,提取空间关系  [batch(batch_size - pre_step), node_num, feature_dim] -> [batch, node_num, temp_size_conv2]
        out = self.conv1(input, edge_index, edge_weight)
        out = self.conv2(out, edge_index, edge_weight)
        # 防止过拟合 ReLU + Dropout
        out = self.relu1(out)
        out = self.dropout(out)
        # 多层感知机,整合各站点的各自特征数据为指定维度的特征张量  [batch, node_num, temp_size_conv2] -> [batch, node_num, gcn_output_size]
        out = self.MLP(out)
        # 置换站点与特征数的维度  [batch, node_num, gcn_output_size] -> [batch, gcn_output_size, node_num]
        out = out.transpose(1, 2)
        # 整合t时刻所有站点数据到一个维度  [batch, gcn_output_size, node_num] -> [batch, gcn_output_size, 1]
        out = self.FC1(out)
        # 将结果降维 [batch, gcn_output_size, 1] -> [batch, gcn_output_size]
        out = torch.squeeze(out, dim=2)

        return out

